from demo_qwen3_api import get_qwen3_client, get_ans_dict_by_search, get_non_empty_keys,\
    get_matched_info_by_jsondict

qwen_client = None
df=None
empty_emb=None
template_json_str=None

def search_and_show(search_txt, threshold=0.5):
    import pandas as pd
    global qwen_client, df,empty_emb,template_json_str
    # 读取csv
    if df is None:
        csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data2600.csv'
        # csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data_api.csv'
        df = pd.read_csv(csv_file_path)
    if template_json_str is None:
        # 读取模板json字符串
        template_json_path = 'demo_qwenvl_clothing.json'
        with open(template_json_path, 'r', encoding='utf-8') as f:
            template_json_str = f.read()
    # 获取qwen3客户端
    if qwen_client is None:
        qwen_client = get_qwen3_client()
    # 获取推理结果
    ans = get_ans_dict_by_search(qwen_client, template_json_str, search_txt)
    import json
    try:
        ans_dict = json.loads(ans)
    except Exception as e:
        print("Failed to parse ans as JSON:", e)
        ans_dict = None
    # 获得非空key链
    non_empty_keys = get_non_empty_keys(ans_dict)
    '''
    key 需要做筛选
    1 只筛选 
        basic_info : season, gender     e.g. ['basic_info','season'] ['basic_info','season']
        design: color.primary color.secondary color.pattern
                category.main category.sub
                sleeve.type sleeve.cuff sleeve.count e.g. ['design', 'sleeve', 'type']...
                collar.type
                pockets
        material: composition
        decoration
        visual
        # 最后的keys只包含以上的key序列
        # 当前 non_empty_keys 中 包含 [ ['design', 'color', 'primary'],[],...  ]
    2 数量 软对比 sim=target_num/num if num >= target_num and num != -1 else 0
    3 正反面 刚性对比
            visual.view_types 中 sim=1 search_type in types else 0
    4 search value 时 target value为空时 加几分？
    '''
    # INSERT_YOUR_CODE

    # ----------- (1) 筛选 key，仅保留指定key链 ------------

    # 定义允许的key链前缀(只判断到某一级即可)
    allowed_prefixes = [
        ['basic_info', 'season'],
        ['basic_info', 'gender'],
        ['design', 'color', 'primary'],
        ['design', 'color', 'secondary'],
        ['design', 'color', 'pattern'],
        ['design', 'category', 'main'],
        ['design', 'category', 'sub'],
        ['design', 'sleeve', 'type'],
        ['design', 'sleeve', 'cuff'],
        ['design', 'sleeve', 'count'],
        ['design', 'collar', 'type'],
        ['design', 'pockets'],         # pockets是整个list,后续get_value_by_keys能拿到
        ['material', 'composition'],
        ['decoration'],                # 只要一级即可，后续靠get_value_by_keys拉取全部
        ['visual'],
    ]

    # 给定实际key链，判定是否属于允许集合
    def is_key_allowed(keys):
        for prefix in allowed_prefixes:
            if len(keys) < len(prefix):
                continue
            # 只要前缀对上即可
            if keys[:len(prefix)] == prefix:
                return True
        return False

    print('[before] non_empty_keys:',non_empty_keys)
    # 对non_empty_keys做一次筛选
    non_empty_keys = [keys for keys in non_empty_keys if is_key_allowed(keys)]
    print('[sifted] non_empty_keys:',non_empty_keys)
    

    # get_matched_info_by_jsondict 这个函数应该全局可用
    from util_for_clipvitb32 import clip_vit_b32, get_embedding
    import torch
    import torch.nn.functional as F
    import heapq
    clip_model = clip_vit_b32()
    if empty_emb is None:
        empty_txt = ''
        empty_emb = get_embedding(clip_model, empty_txt)    
    search_joined_values_list = get_matched_info_by_jsondict(ans_dict, non_empty_keys)
    search_joined_values_embs = [get_embedding(clip_model, item) 
                                    for item in search_joined_values_list]
    search_joined_values = '&&'.join(search_joined_values_list)                                    

    topk = 10
    heap = []

    for idx, row in df.iterrows():
        target_json_dict_str = row.get('desc_json', None)
        img_path = row.get('img_path', '')

        try:
            target_json_dict = json.loads(target_json_dict_str)
        except Exception as e:
            # print("template_json_str 解析失败:", e)
            target_json_dict = {}

        joined_values_list = get_matched_info_by_jsondict(target_json_dict, non_empty_keys)
        
        # 对 每一个 item 都做 相似度，每个相似度 when if < thre，则continue
        # 否则 相似度累加 sum_sim , 按照这个排序
        sim_list = []
        satisfy = True
        for joined_value, search_joined_value, search_emb in \
            zip(joined_values_list, search_joined_values_list, search_joined_values_embs):
            # 由于clip对数量不敏感, 如果 search_joined_value 可以转换为 数字类型
            # 则 尝试将 joined_value 也进行转换,如果 转换失败 或者 数字不等,直接satisfy = False and break
            # 判断search_joined_value是否可转为数字，且不为-1
            if search_joined_value.isdigit() or (search_joined_value.startswith('-') and search_joined_value[1:].isdigit()):
                search_num = int(search_joined_value)

                # key中不会存在 num =-1 的情况, 已被筛除

                if search_num == 0:
                    if joined_value == "": # len(pockets)==0 由于没找到index因此返回None=>''
                        sim_list.append(1)
                        continue
                    # 检索 要求 没有, 则 target 也得 没有,此时target 必须为 -1
                    # else satisfy=False
                    if joined_value.isdigit() or (joined_value.startswith('-') and joined_value[1:].isdigit()):
                        joined_num = int(joined_value)
                        # 检索要求没有，目标只能是-1或0
                        if joined_num in [0, -1]:
                            sim_list.append(1)
                            continue
                        else:
                            satisfy = False
                            break
                    else:
                        satisfy = False
                        break
                # search num != 0
                if joined_value.isdigit() or (joined_value.startswith('-') and joined_value[1:].isdigit()):
                    joined_num = int(joined_value)

                    # sim = search_num / joined_num if joined_num >= search_num and joined_num != -1 else 0
                    sim = 1 if search_num == joined_num else 0
                    # if sim == 0:
                    if sim < threshold:
                        satisfy = False
                        break
                    sim_list.append(sim)
                    continue
                else:
                    satisfy = False
                    break

            # 针对正反面强约束，如果search类型为正反面，只允许visual.view_types完全匹配（而不是用CLIP embedding比）
            # 假设non_empty_keys, joined_values_list, search_joined_values_list, search_joined_values_embs 位置一一对应
            # 若某key是visual.view_types，严格判定
            cur_keys = non_empty_keys[joined_values_list.index(joined_value)] if joined_value in joined_values_list else []
            # 判断key链是否包含visual和view_types
            if len(cur_keys) >= 2 and cur_keys[-2:] == ["visual", "view_types"]:
                # search_type只可能为这些
                types = {"front", "back", "side", "detail"}
                # 严格正反面
                # search_joined_value和joined_value都可能为'front', 'back', 'side', 'detail' 或用逗号拼接
                # 只要search_joined_value里的type都在joined_value里，才算satisfy, 否则sim=0，break
                sv_set = set([v.strip() for v in search_joined_value.split(",")])
                jv_set = set([v.strip() for v in joined_value.split(",")])
                if not sv_set.issubset(jv_set):
                    satisfy = False
                    break
                # 逻辑过了，给sim=1
                sim_list.append(1)
                continue

            
            joined_values_emb = get_embedding(clip_model, joined_value)

            # base_sim = F.cosine_similarity(
            #     torch.tensor(empty_emb).unsqueeze(0),
            #     torch.tensor(joined_values_emb).unsqueeze(0)
            # ).item()
            ## test cosim
            search_sim = F.cosine_similarity(
                torch.tensor(search_emb).unsqueeze(0),
                torch.tensor(joined_values_emb).unsqueeze(0)
            ).item()
            ## test L1 sim
            # 使用欧氏距离作为相似度
            # 欧氏距离越小越相似，因此这里sim可以用 1 / (1 + 距离)
            # euclidean_dist = torch.norm(torch.tensor(search_emb) - torch.tensor(joined_values_emb), p=2).item()
            # search_sim = 1.0 / (1.0 + euclidean_dist)

            if search_sim < threshold:
                satisfy = False
                break

            # relative_sim = search_sim - base_sim
            sim_list.append( search_sim )
        
        if satisfy:
            sum_sim = sum(sim_list)
        # if relative_sim >= 0:
            # item = (relative_sim, search_sim, base_sim, idx, img_path, joined_values)
            item = ( sum_sim , [f'{sim:.4f}' for sim in sim_list], 
                    idx, img_path, '&&'.join(joined_values_list))
            if len(heap) < topk:
                heapq.heappush(heap, item)
            else:
                if sum_sim > heap[0][0]:
                    heapq.heappushpop(heap, item)


    # 排序，返回info和图片路径
    heap_sorted = sorted(heap, key=lambda x: -x[0])
    res_list = []
    img_list = []
    from utils.util_flux import process_img_1024
    for sim, sim_list, idx, path, tar_value_str in heap_sorted:
        res_list.append( f"TOPK sum_sim={sim:.4f} search_sim={sim_list}  \
                         idx={idx} path= {path} desc={tar_value_str}")
        img_list.append(
            process_img_1024(path)
        )

    # 返回search_txt, 图片列表，内容列表
    # non_empty_keys search_joined_values_list
    search_txt_list = [f'{k}:{v}' for k,v in zip(non_empty_keys, search_joined_values_list)]
    search_txt = '\n'.join(search_txt_list)
    return search_txt, img_list, '\n'.join(res_list)

def start_gradio(port):
    import gradio as gr
    with gr.Blocks() as demo:
        gr.Markdown("# QwenVL3 Gradio Demo (以文搜图)")
        with gr.Row():
            # with gr.Column():
            search_box = gr.Textbox(label="输入搜索描述", placeholder="如：春秋两季黑色大衣，两个口袋", lines=2)
        with gr.Row():
            threshold_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.01, label="阈值 threshold", interactive=True)
            submit_btn = gr.Button("开始搜索")
            # with gr.Column():
        with gr.Row():
            output_search = gr.Textbox(label="搜索内容", interactive=False)
        with gr.Row():
            gallery = gr.Gallery(label="TopK图片预览", columns=5, height="auto")
        with gr.Row():
            output_list = gr.Textbox(label="TopK内容", lines=10, interactive=False)
        submit_btn.click(
            search_and_show,
            inputs=[search_box, threshold_slider],
            outputs=[output_search, gallery, output_list]
        )
    demo.launch(server_name="0.0.0.0", server_port=port)

# 若本py被直接调用则自动启动界面
if __name__ == '__main__':
    import os
    os.environ['GRADIO_TEMP_DIR'] = '/mnt/nas/shengjie/tmp'
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
    parser.add_argument('-p', '--port', type=int, default=20023, help='CUDA device id')
    args, unknown = parser.parse_known_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

    start_gradio(args.port)